import torch
from modelscope.utils.constant import Tasks
from modelscope.pipelines import pipeline
from modelscope.preprocessors.image import load_image

pipeline = pipeline(task=Tasks.multi_modal_embedding, model='damo/multi-modal_clip-vit-large-patch14_336_zh', model_revision='v1.0.1', device='cpu')
input_img = load_image('https://p1.a.yximgs.com/upic/2023/12/27/13/BMjAyMzEyMjcxMzAwMzdfMjkwOTEwNDc5N18xMjA2NTYyODIyNDRfMl8z_Be33b730fcb8aaf000ba5e58bf8c3681e.jpg?tag=1-1708460644-unknown-0-huswbkiuiq-0ac88680ace16101&clientCacheKey=3x9ghy3dj8y2z5a.jpg&di=70208a94&bp=14764') # 支持皮卡丘示例图片路径/本地图片 返回PIL.Image
input_texts = ["小火龙", "白色衬衫"]

# 支持一张图片(PIL.Image)或多张图片(List[PIL.Image])输入，输出归一化特征向量
img_embedding = pipeline.forward({'img': input_img})['img_embedding'] # 2D Tensor, [图片数, 特征维度]

# 支持一条文本(str)或多条文本(List[str])输入，输出归一化特征向量
text_embedding = pipeline.forward({'text': input_texts})['text_embedding'] # 2D Tensor, [文本数, 特征维度]

# 计算图文相似度
with torch.no_grad():
    # 计算内积得到logit，考虑模型temperature
    logits_per_image = (img_embedding / pipeline.model.temperature) @ text_embedding.t()
    # 根据logit计算概率分布
    probs = logits_per_image.softmax(dim=-1).cpu().numpy()

# similarityNum = probs[0][1] * 100
print("图文匹配概率:", probs[0])
# result = True if similarityNum >= 85 else False
# print("图文匹配概率:", result)